import argparse
import re
import os
import time
import numpy as np
import math
from collections import OrderedDict
import torch
from utils import random_seed, create_result_dir, Logger, TableLogger, AverageMeter
from attack import AttackPGD
from adamw import AdamW
from model.model import Model, set_eps, get_eps
from model.norm_dist import set_p_norm, get_p_norm
import matplotlib.pyplot as plt
from torchvision import transforms

parser = argparse.ArgumentParser(description='Adversarial Robustness')
parser.add_argument('--dataset', default='MNIST', type=str)
parser.add_argument('--model', default='MLPFeature(depth=4,width=4)', type=str) #mlp,conv
parser.add_argument('--predictor-hidden-size', default=512, type=int) # 0 means not to use linear predictor
parser.add_argument('--loss', default='cross_entropy', type=str) #cross_entropy, hinge

parser.add_argument('--p-start', default=8.0, type=float)
parser.add_argument('--p-end', default=1000.0, type=float)
parser.add_argument('--kappa', default=1.0, type=float)
parser.add_argument('--epochs', default='0,50,50,350,400', type=str) # epoch1-epoch3: inc eps; epoch2-epoch4: inc p

parser.add_argument('--eps-train', default=None, type=float)
parser.add_argument('--eps-test', default=None, type=float)

parser.add_argument('-b', '--batch-size', default=256, type=int)
parser.add_argument('--lr', default=0.01, type=float)
parser.add_argument('--beta1', default=0.9, type=float)
parser.add_argument('--beta2', default=0.99, type=float)
parser.add_argument('--epsilon', default=1e-10, type=float)
parser.add_argument('--wd', default=0.0, type=float)

parser.add_argument('--start-epoch', default=0, type=int)
parser.add_argument('--checkpoint', default=None, type=str)

parser.add_argument('--gpu', default=-1, type=int, help='GPU id to use')
parser.add_argument('--dist-url', default='tcp://localhost:23456')
parser.add_argument('--world-size', default=1)
parser.add_argument('--rank', default=0)

parser.add_argument('-p', '--print-freq', default=50, type=int, metavar='N', help='print frequency')
parser.add_argument('--result-dir', default='result/', type=str)
parser.add_argument('--filter-name', default='', type=str)
parser.add_argument('--seed', default=2020, type=int)
parser.add_argument('--visualize', action='store_true')

args = parser.parse_args()

def parse_function_call(s):
    s = re.split(r'[()]', s)
    if len(s) == 1:
        return s[0], {}
    name, params, _ = s
    params = re.split(r',\s*', params)
    params = dict([p.split('=') for p in params])
    for key, value in params.items():
        try:
            params[key] = int(params[key])
        except ValueError:
            try:
                params[key] = float(params[key])
            except ValueError:
                pass
    return name, params

'''
load data
'''
from dataset import load_data, get_statistics, default_eps, input_dim
train_loader, test_loader = load_data(args.dataset, 'data/', args.batch_size // args.world_size, False, augmentation=False, classes=None)
mean, std = get_statistics(args.dataset)
num_classes = len(train_loader.dataset.classes)

up = torch.FloatTensor((1 - mean) / std).view(-1, 1, 1).cuda(0)
down = torch.FloatTensor((0 - mean) / std).view(-1, 1, 1).cuda(0)

'''
load model
'''
from model.bound_module import Predictor, BoundFinalIdentity
from model.mlp import MLPFeature, MLP, Denoising, DenoisingAdvance, Denoising_MLP
from model.conv import ConvFeature, Conv, Denoising_Conv, Denoising_ConvFeature
from model.resnet import ResNetDenoiseModel
model_name, params = parse_function_call(args.model)
for k, v in params.items():
    if isinstance(v, str) and (locals().get(v) is not None):
        params[k] = locals()[v]
if args.predictor_hidden_size > 0:
    model = locals()[model_name](input_dim=input_dim[args.dataset], **params)
    predictor = Predictor(model.out_features, args.predictor_hidden_size, num_classes)
else:
    model = locals()[model_name](input_dim=input_dim[args.dataset], num_classes=num_classes, **params)
    predictor = BoundFinalIdentity()
model = Model(model, predictor, eps=0)
model = model.cuda(0)
model.load_state_dict(torch.load('result/CIFAR10_Denoising_ConvFeature(width=1,denoising_module=DenoisingAdvance)_hid256_p8.0_p_end1000.0_kappa0.5_epochs0,60,60,450,500_eps_train0.35_eps_test0.3_bs512_lr0.02_beta10.9_wd0.005__2020__0/model.pth')['state_dict'], False)
from model.norm_dist import set_p_norm
set_p_norm(model, float('inf'))

'''
generate adv examples
'''
attacker = AttackPGD(model, args.eps_test, step_size=args.eps_test / 4, num_steps=20, up=up, down=down)
for id, (inputs, targets) in enumerate(test_loader):
    inputs = inputs.cuda(0, non_blocking=True)
    targets = targets.cuda(0, non_blocking=True)
    perturb = attacker.find(inputs, targets, random_start=bool(torch.rand(1) < .8))
    for i in range(5):
        p = perturb[i].cpu() * std
        p[0] += mean[0]
        p[1] += mean[1]
        p[2] += mean[2]
        torch.save(p, f'get_example_{i}_adv.txt')
    for i in range(5):
        p = inputs[i].cpu() * std
        p[0] += mean[0]
        p[1] += mean[1]
        p[2] += mean[2]
        torch.save(p, f'get_example_{i}.txt')
        print(targets[i])
    break